from openai import OpenAI
from flask import current_app
from models.database import Settings
import json

class AIService:
    """AI服务类"""
    
    def __init__(self, api_key=None, model=None):
        """初始化AI服务
        
        Args:
            api_key: API密钥
            model: 模型名称
        """
        self.api_key = api_key
        self.model = model or "deepseek-chat"
        self.client = None
        
    def init_client(self):
        """初始化OpenAI客户端"""
        if not self.api_key:
            # 从数据库获取设置
            settings = Settings.query.first()
            if settings:
                self.api_key = settings.api_key
                self.model = settings.api_model or self.model
                
        if not self.api_key:
            raise ValueError("API key is required")
            
        self.client = OpenAI(
            api_key=self.api_key,
            base_url="https://api.deepseek.com"
        )
        
    def chat(self, messages, stream=False, **kwargs):
        """发送聊天请求
        
        Args:
            messages: 消息列表
            stream: 是否使用流式响应
            **kwargs: 其他参数
            
        Returns:
            聊天响应
        """
        if not self.client:
            self.init_client()
            
        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            stream=stream,
            **kwargs
        )
        
        if stream:
            return response
        else:
            return {
                'content': response.choices[0].message.content,
                'role': response.choices[0].message.role,
                'finish_reason': response.choices[0].finish_reason
            }
            
    def stream_handler(self, response):
        """处理流式响应
        
        Args:
            response: 流式响应对象
            
        Yields:
            每个响应片段
        """
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                yield {
                    'content': chunk.choices[0].delta.content,
                    'role': chunk.choices[0].delta.role if hasattr(chunk.choices[0].delta, 'role') else None,
                    'finish_reason': chunk.choices[0].finish_reason
                } 